package cn.iocoder.yudao.framework.datasource.processor;

import cn.iocoder.yudao.framework.datasource.local.ThreadLocalContext;
import cn.iocoder.yudao.framework.msdb.ICustomizeRoutingDataSource;
import com.baomidou.dynamic.datasource.processor.DsProcessor;
import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import org.aopalliance.intercept.MethodInvocation;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.util.StringUtils;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.util.Objects;

/**
 * TODO
 *
 * @className: DsRestControllerProcessor
 * @author: tangkh
 * @date: 2023-09-19 17:57
 */
public class DsRestControllerProcessor extends DsProcessor {
    private static final Logger log = LoggerFactory.getLogger(DsRestControllerProcessor.class);
    @Value("${spring.datasource.dynamic.extends.dbkey:}")
    private String dbKey;
    @Value("${spring.datasource.dynamic.extends.multi-tenant.enable:false}")
    private boolean multiTenantEnable;
    public static final String SLAVE = "slave";
    private static final String READ_ONLY_REGEX = "[\\s\\S]*.*#query=ReAdOnLY#.*[\\s\\S]*";
    @Value("${spring.datasource.dynamic.primary}")
    private String primary;
    @Value("${spring.datasource.dynamic.extends.dbkey:}")
    private String defaultDbkey;
    @Autowired(
            required = false
    )
    private ICustomizeRoutingDataSource contextPathInterceptor;

    public DsRestControllerProcessor() {
    }

    @Override
    public boolean matches(String key) {
        return true;
    }

    @Override
    public String doDetermineDatasource(MethodInvocation invocation, String key) {
        return this.getMaster(key);
    }

    private String getDatasourceName(String key, String orgId) {
        if (!Objects.isNull(orgId)) {
            String substring = "".equals(key) ? this.dbKey : key;
            return orgId + "_" + substring;
        } else {
            return StringUtils.isEmpty(key) ? this.dbKey : key;
        }
    }

    private String getMaster(String key) {
        String master = StringUtils.hasText(key) ? key : DynamicDataSourceContextHolder.peek();
        String dbKey = (String)ThreadLocalContext.getAttachment("dbKey");
        if (StringUtils.hasText(dbKey)) {
            master = dbKey;
        } else if (this.contextPathInterceptor != null) {
            String value = this.contextPathInterceptor.getDbKey();
            if (StringUtils.hasText(value)) {
                master = value;
            } else {
                master = StringUtils.hasText(master) ? master : this.defaultDbkey;
            }
        }

        if (this.multiTenantEnable) {
            RequestAttributes requestAttributes = RequestContextHolder.getRequestAttributes();
            if (requestAttributes != null) {
                HttpServletRequest request = ((ServletRequestAttributes)requestAttributes).getRequest();
                String orgId = request.getParameter("gwOrgId");
                master = this.getDatasourceName(master, orgId);
            } else {
                String orgId = (String) ThreadLocalContext.getAttachment("orgId");
                master = this.getDatasourceName(master, orgId);
            }
        }

        if (!StringUtils.hasText(master)) {
            master = this.defaultDbkey;
        }

        return master;
    }
}
